# 人物小百科
from flask import Flask, request, jsonify
import requests
from langchain.prompts import PromptTemplate
from langchain.llms import OpenAI
import os
from flask_cors import CORS
import re
import json

app = Flask(__name__)
CORS(app)  # 添加跨域支持

# 设置 OpenAI API 密钥
os.environ["OPENAI_API_KEY"] = "sk-fhfopxxokpblqjbivoyhwhftlbjwrgeddccbytodqipkaqum"
base_url="https://api.siliconflow.cn/v1"
# 配置 SiliconFlow 的 OpenAI 兼容 API
llm = OpenAI(
    openai_api_base=base_url,
    api_key=os.environ.get("OPENAI_API_KEY"),
    # 不是推理模型
    #model_name="THUDM/glm-4-9b-chat",
    model_name="THUDM/GLM-4-32B-0414",
    temperature=0.7
)
# 定义 Prompt 模板
prompt_template = """请根据历史知识，提供关于 {character_name} 的信息。
如果您知道相关信息，请只按照以下格式强行输出：
{{
    "name": "{character_name}",
    "bio": "人物简介",
    "funStory": "趣味故事",
    "keyEvents": ["关键事件1", "关键事件2"],
    "relatedCharacters": ["关联人物1", "关联人物2"]
}}
关联人物不要重复.
不要回答markdown格式，不要回答其他格式。
如果您不知道人物的相关信息，请直接回复："我不知道关于 {character_name} 的信息。"
"""

prompt = PromptTemplate(
    input_variables=["character_name"],
    template=prompt_template,
)
def remove_markdown(result):
    # 移除常见的 Markdown 语法元素，如标题、列表、链接、强调等
    result = re.sub(r'#+\s', '', result)  # 移除标题
    result = re.sub(r'- \[.*?\]\(.*?\)', '', result)  # 移除带链接的列表项
    result = re.sub(r'- ', '', result)  # 移除无序列表项
    result = re.sub(r'\[.*?\]\(.*?\)', '', result)  # 移除链接
    result = re.sub(r'\*\*.*?\*\*', lambda x: x.group(0).replace('**', ''), result)  # 移除加粗
    result = re.sub(r'\*.*?\*', lambda x: x.group(0).replace('*', ''), result)  # 移除斜体
    # 移除代码块，同时处理代码块格式中的 json、python 等字符
    result = re.sub(r'`(?:json|python)?(.*?)`', lambda x: x.group(1), result)  # 移除代码块及可能的格式标识
    return result

def get_character_info():
    character_name = request.args.get('character_name')
    print(character_name)
    print(character_name)
    print(character_name)
    print(character_name)
    print(character_name)
    if not character_name:
        return jsonify({'error': '请提供人物姓名'}), 400

    # 格式化 Prompt
    formatted_prompt = prompt.format(character_name=character_name)

    try:
        # 调用大模型获取信息
        result = llm(formatted_prompt)
        print(result)
        print('--------------------------------')
        # 在调用大模型获取信息后调用此函数处理结果
        result = remove_markdown(result)
        print('--------------------------------')
        try:
            import json
            # 尝试找到第一个有效的 JSON 内容
            start_index = result.find('{')
            end_index = result.rfind('}')
            if start_index != -1 and end_index != -1:
                valid_json_str = result[start_index:end_index + 1]
                character_info = json.loads(valid_json_str)
                return jsonify({'data': character_info})
            else:
                print("无法解析为 JSON，返回原始结果")
                return jsonify({'data': result})
        except json.JSONDecodeError:
            print("无法解析为 JSON，返回原始结果")
            return jsonify({'data': result})
    except Exception as e:
        return jsonify({'error': str(e)}), 500
        
if __name__ == '__main__':
    app.run(debug=True, use_reloader=True, port=5001)